/**
    Copyright (c) 2010-2011 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "cutil_math.h"

/**
    @file   
    @brief  Mandelbrot cuda program
    @date   2010/02/27
*/

struct ConstantBuf
{
     // bounds of the complex number plane
    float fMinRe;
    float fMaxRe;
    float fMinIm;
    float fMaxIm;
    
    // step size for the complex number plane
    float fStepRe;
    float fStepIm;

    // width*height of the data buffer
    int        nSize;

    // width of the data buffer
    int        nWidth;

    // maximum iterations to decide whether an complex number is in the mandelbrot set.
    int        nMaxIterations;

    int  padding[ 3 ];        // constant buffer ByteWidth has to be multiple of 16
};

__constant__ ConstantBuf cb;

__global__ void cuda_kernel_mandelbrot(unsigned char* surface, int pitch )
{
    const int x = blockIdx.x*blockDim.x + threadIdx.x;
    const int y = blockIdx.y*blockDim.y + threadIdx.y;

    int nIteration=0;
            
    float4 color = make_float4( 0.0f, 0.0f, 0.0f, 1.0f );
    
    float cIm = cb.fMaxIm - y*cb.fStepIm;
    float cRe = cb.fMinRe + x*cb.fStepRe;

    float zRe = cRe;
    float zIm = cIm;
            
    for( nIteration=0; nIteration<cb.nMaxIterations; ++nIteration )
    {
        float zRe2 = zRe*zRe;
        float zIm2 = zIm*zIm;
                
        /// check whether absolute value of z has gotten bigger than 2
        if( zRe2 + zIm2 > 4 ) {
            break;
        }

        // calculate next z(n)
        zIm = 2*zRe*zIm + cIm;
        zRe = zRe2-zIm2 + cRe;
    }
                        
    if( nIteration != cb.nMaxIterations ) { 
        if( nIteration < cb.nMaxIterations/2 ) {
            color = make_float4( nIteration / ((float)( cb.nMaxIterations )/2.0f), 0.0f, 0.0f, 1.0f );
        }
        else {
            float fLerp = ( cb.nMaxIterations-nIteration ) / ((float)( cb.nMaxIterations )/2.0f);
            color = lerp( make_float4( 1.0f, 0.0f, 0.0f, 1.0f ),make_float4( 1.0f, 1.0f, 1.0f, 1.0f ), fLerp );
        }    
    }

    if( x+cb.nWidth*y < cb.nSize && x < cb.nWidth )
    {
        float* pixel;
        pixel = (float*)(surface + y*pitch) + 4*x;
        pixel[ 0 ] = color.x;
        pixel[ 1 ] = color.y;
        pixel[ 2 ] = color.z;
        pixel[ 3 ] = color.w;
    }
}

extern "C" 
void cuda_mandelbrot( void* surface, int width, int height, int pitch )
{
    dim3 threads = dim3( 16, 16 );
    dim3 blocks = dim3( (width+threads.x-1)/threads.x, (height+threads.y-1)/threads.y );

    cuda_kernel_mandelbrot<<<blocks,threads>>>( (unsigned char*)surface, pitch );
}
